package cn.itcast.rbrule;

import com.netflix.client.config.CommonClientConfigKey;
import com.netflix.client.config.IClientConfig;
import com.netflix.client.config.IClientConfigKey;
import com.netflix.loadbalancer.*;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @Author: zhaojie
 * @Description: 自定义权重负载均衡，根据配置的权重数，数字越大被访问的概率越高，如果未获取到配置，默认使用随机负载策略
 * 新增服务掉线、上线不同场景加载逻辑
 */
public class WeightedRule extends RoundRobinRule {
    public WeightedRule() {

    }

    private Logger logger = LoggerFactory.getLogger(WeightedRule.class);

    //权重区间
    private volatile List<Double> accumulatedWeights = new ArrayList();

    //权重配置
    private volatile Map<String, Double> map = new LinkedHashMap<>();

    //恒定规则
    private Map<String, Double> finalMap = new LinkedHashMap<>();

    //定时任务间隔
    private int serverWeightTaskTimerInterval = 30000;
    //总份额

    private double totalResponseTime = 100D;

    //定时器
    protected Timer serverWeightTimer = null;

    protected AtomicBoolean serverWeightAssignmentInProgress = new AtomicBoolean(false);

    protected AtomicBoolean resetFlag = new AtomicBoolean(false);
    //服务名
    String name = "unknown";

    //自定义config
    public static final IClientConfigKey<String> weightRule = new CommonClientConfigKey<String>("weightRule") {
    };


    public void setLoadBalancer(ILoadBalancer lb) {
        super.setLoadBalancer(lb);
        if (lb instanceof BaseLoadBalancer) {
            this.name = ((BaseLoadBalancer) lb).getName();
        }
        this.initialize(lb);
    }

    //初始化权重
    void initialize(ILoadBalancer lb) {
        if (this.serverWeightTimer != null) {
            this.serverWeightTimer.cancel();
        }
        this.serverWeightTimer = new Timer("NFLoadBalancer-serverWeightTimer-" + this.name, true);
        this.serverWeightTimer.schedule(new WeightedRule.DynamicServerWeightTask(), 0L, (long) this.serverWeightTaskTimerInterval);
        WeightedRule.ServerWeight sw = new WeightedRule.ServerWeight();
        sw.maintainWeights();
        Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {
            public void run() {
                logger.info("Stopping NFLoadBalancer-serverWeightTimer-" + WeightedRule.this.name);
                WeightedRule.this.serverWeightTimer.cancel();
            }
        }));
    }

    /**
     * 定时任务获取权重
     */
    class DynamicServerWeightTask extends TimerTask {
        DynamicServerWeightTask() {
        }

        public void run() {
            WeightedRule.ServerWeight serverWeight = WeightedRule.this.new ServerWeight();

            try {
                serverWeight.maintainWeights();
            } catch (Exception var3) {
                logger.error("Error running DynamicServerWeightTask for {}", WeightedRule.this.name, var3);
            }

        }
    }

    //如果服务全部恢复 则还是采用配置文件权重配置
    void setYMlWeightRule(List<Server> allServers) {
        if (WeightedRule.this.resetFlag.compareAndSet(false, true)) {
            if (allServers.size() == finalMap.size()) {
                //判断所启动服务和所配置服务是否一致
                Set<String> strings = finalMap.keySet();
                int i = 0;
                for (Server allServer : allServers) {
                    String key = allServer.getHost() + ":" + allServer.getPort();
                    if (strings.contains(key)) {
                        i = i + 1;
                    }
                }
                if (i == allServers.size()) {
                    this.map = finalMap;
                }
            }
        }
    }

    class ServerWeight {
        ServerWeight() {
        }

        public void maintainWeights() {
            ILoadBalancer lb = WeightedRule.this.getLoadBalancer();
            if (lb != null) {
                if (WeightedRule.this.serverWeightAssignmentInProgress.compareAndSet(false, true)) {
                    try {
                        logger.info("Weight adjusting job started");
                        AbstractLoadBalancer nlb = (AbstractLoadBalancer) lb;
                        LoadBalancerStats stats = nlb.getLoadBalancerStats();
                        if (stats != null) {
                            List<Server> allServers1 = nlb.getAllServers();

                            //如果服务全部恢复 则还是采用配置文件权重配置
                            setYMlWeightRule(allServers1);

                            //由于可能服务掉线，map需要重新设置
                            if (map.size() > allServers1.size()) {
                                map = new LinkedHashMap<>();
                                WeightedRule.this.resetFlag.set(false);
                            }

                            //由于可能后续添加服务 所以重新设置权重 采用平均分
                            while (map == null || map.size() < allServers1.size()) {
                                List<Server> allServers = nlb.getAllServers();
                                Map<String, Double> newMap = new LinkedHashMap<>();
                                for (Server allServer : allServers) {
                                    String serviceKey = allServer.getHost() + ":" + allServer.getPort();
                                    newMap.put(serviceKey, totalResponseTime / nlb.getAllServers().size());
                                }
                                map = newMap;
                                WeightedRule.this.resetFlag.set(false);
                            }

                            //权重累加
                            Double weightSoFar = 0.0D;

                            //权重分配list
                            List<Double> finalWeights = new ArrayList();
                            Iterator var20 = nlb.getAllServers().iterator();

                            while (var20.hasNext()) {
                                Server serverx = (Server) var20.next();
                                //当前权重分区
                                double weight = map.get(serverx.getHost() + ":" + serverx.getPort());
                                weightSoFar = weightSoFar + weight;
                                finalWeights.add(weightSoFar);
                            }
                            WeightedRule.this.setWeights(finalWeights);
                            return;
                        }
                    } catch (Exception var16) {
                        logger.error("Error calculating server weights", var16);
                        return;
                    } finally {
                        WeightedRule.this.serverWeightAssignmentInProgress.set(false);
                    }

                }
            }
        }

    }

    void setWeights(List<Double> weights) {
        this.accumulatedWeights = weights;
    }

    private final Random random = new Random();

    public void setWeightedRule(List<Server> allList) {
        Map<String, Double> newMap = new LinkedHashMap<>();
        for (Server server : allList) {
            String serviceKey = server.getHost() + ":" + server.getPort();
            newMap.put(serviceKey, totalResponseTime / allList.size());
        }
        this.map = newMap;
    }

    public void setWeight(List<Server> allList) {
        //权重累加
        Double weightSoFar = 0.0D;
        //权重分配list
        List<Double> finalWeights = new ArrayList();
        for (Server server : allList) {
            //100份 分成若干份
            //当前权重
            double weight = map.get(server.getHost() + ":" + server.getPort());
            weightSoFar = weightSoFar + weight;
            finalWeights.add(weightSoFar);
        }
        WeightedRule.this.setWeights(finalWeights);
    }

    public Server choose(ILoadBalancer lb, Object key) {
        if (lb == null) {
            return null;
        } else {
            Server server = null;
            while (server == null) {
                List<Double> currentWeights = this.accumulatedWeights;
                if (Thread.interrupted()) {
                    return null;
                }
                List<Server> allList = lb.getAllServers();

                int serverCount = allList.size();
                if (serverCount == 0) {
                    return null;
                }

                int serverIndex = 0;
                double maxTotalWeight = currentWeights.size() == 0 ? 0.0D : (Double) currentWeights.get(currentWeights.size() - 1);
                if (maxTotalWeight < 0.001D) {
                    server = super.choose(this.getLoadBalancer(), key);
                    if (server == null) {
                        return server;
                    }
                } else {
                    double randomWeight = this.random.nextDouble() * maxTotalWeight;
                    int n = 0;

                    for (Iterator var13 = currentWeights.iterator(); var13.hasNext(); ++n) {
                        Double d = (Double) var13.next();
                        if (d >= randomWeight) {
                            serverIndex = n;
                            break;
                        }
                    }
                    server = (Server) allList.get(serverIndex);
                }

                if (server == null) {
                    Thread.yield();
                } else {
                    if (server.isAlive()) {
                        return server;
                    }

                    server = null;
                }
            }

            return server;
        }


    }

    //获取配置
    public void initWithNiwsConfig(IClientConfig clientConfig) {
        String ruleClassName = (String) clientConfig.getProperty(weightRule);
        if (StringUtils.isNotEmpty(ruleClassName)) {
            if (this.map.size() == 0) {
                String[] split = ruleClassName.split(",");
                for (String sp : split) {
                    String[] split1 = sp.split("-");
                    this.map.put(split1[0], Double.valueOf(split1[1]));
                }
                this.finalMap = map;
            }
        }
    }
}
